{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "google",
   "metadata": {},
   "source": [
    "##### Copyright 2025 Google LLC."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "apache",
   "metadata": {},
   "source": [
    "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "you may not use this file except in compliance with the License.\n",
    "You may obtain a copy of the License at\n",
    "\n",
    "    http://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "See the License for the specific language governing permissions and\n",
    "limitations under the License.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "basename",
   "metadata": {},
   "source": [
    "# test_scheduling_sat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "link",
   "metadata": {},
   "source": [
    "<table align=\"left\">\n",
    "<td>\n",
    "<a href=\"https://colab.research.google.com/github/google/or-tools/blob/main/examples/notebook/examples/test_scheduling_sat.ipynb\"><img src=\"https://raw.githubusercontent.com/google/or-tools/main/tools/colab_32px.png\"/>Run in Google Colab</a>\n",
    "</td>\n",
    "<td>\n",
    "<a href=\"https://github.com/google/or-tools/blob/main/examples/python/test_scheduling_sat.py\"><img src=\"https://raw.githubusercontent.com/google/or-tools/main/tools/github_32px.png\"/>View source on GitHub</a>\n",
    "</td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "doc",
   "metadata": {},
   "source": [
    "First, you must install [ortools](https://pypi.org/project/ortools/) package in this colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "install",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install ortools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "description",
   "metadata": {},
   "source": [
    "\n",
    "Solves a test scheduling problem.\n",
    "\n",
    "Tests must be run by an operator. Tests have a duration and a power consumption.\n",
    "\n",
    "Operators draw power from power supplies. The mapping between operators and\n",
    "power supplies is given.\n",
    "\n",
    "Power supplies have a maximum power they can deliver.\n",
    "\n",
    "Can we schedule the tests so that the power consumption of each power supply is\n",
    "always below its maximum power, and the total makespan is minimized?\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "code",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections.abc import Sequence\n",
    "import io\n",
    "from typing import Dict, Tuple\n",
    "\n",
    "from ortools.sat.colab import flags\n",
    "import pandas as pd\n",
    "\n",
    "from google.protobuf import text_format\n",
    "from ortools.sat.python import cp_model\n",
    "\n",
    "\n",
    "_PARAMS = flags.define_string(\n",
    "    \"params\",\n",
    "    \"num_search_workers:16,log_search_progress:true,max_time_in_seconds:45\",\n",
    "    \"Sat solver parameters.\",\n",
    ")\n",
    "\n",
    "\n",
    "def build_data() -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:\n",
    "    \"\"\"Build the data frame.\"\"\"\n",
    "    tests_str = \"\"\"\n",
    "  Name Operator    TestTime    AveragePower\n",
    "   T1     O1           300            200\n",
    "   T2     O1           150             40\n",
    "   T3     O2           100             65\n",
    "   T4     O2           250            150\n",
    "   T5     O3           210            140\n",
    "  \"\"\"\n",
    "\n",
    "    operators_str = \"\"\"\n",
    "  Operator Supply\n",
    "      O1      S1\n",
    "      O2      S2\n",
    "      O3      S2\n",
    "  \"\"\"\n",
    "\n",
    "    supplies_str = \"\"\"\n",
    "  Supply  MaxAllowedPower\n",
    "   S1        230\n",
    "   S2        210\n",
    "  \"\"\"\n",
    "\n",
    "    tests_data = pd.read_table(io.StringIO(tests_str), sep=r\"\\s+\")\n",
    "    operators_data = pd.read_table(io.StringIO(operators_str), sep=r\"\\s+\")\n",
    "    supplies_data = pd.read_table(io.StringIO(supplies_str), sep=r\"\\s+\")\n",
    "\n",
    "    return (tests_data, operators_data, supplies_data)\n",
    "\n",
    "\n",
    "def solve(\n",
    "    tests_data: pd.DataFrame,\n",
    "    operator_data: pd.DataFrame,\n",
    "    supplies_data: pd.DataFrame,\n",
    ") -> None:\n",
    "    \"\"\"Solve the scheduling of tests problem.\"\"\"\n",
    "\n",
    "    # Parses data.\n",
    "    operator_to_supply: Dict[str, str] = {}\n",
    "    for _, row in operator_data.iterrows():\n",
    "        operator_to_supply[row[\"Operator\"]] = row[\"Supply\"]\n",
    "\n",
    "    supply_to_max_power: Dict[str, int] = {}\n",
    "    for _, row in supplies_data.iterrows():\n",
    "        supply_to_max_power[row[\"Supply\"]] = row[\"MaxAllowedPower\"]\n",
    "\n",
    "    horizon = tests_data[\"TestTime\"].sum()\n",
    "\n",
    "    # OR-Tools model.\n",
    "    model = cp_model.CpModel()\n",
    "\n",
    "    # Create containers.\n",
    "    tests_per_supply: Dict[str, Tuple[list[cp_model.IntervalVar], list[int]]] = {}\n",
    "    test_supply: Dict[str, str] = {}\n",
    "    test_starts: Dict[str, cp_model.IntVar] = {}\n",
    "    test_durations: Dict[str, int] = {}\n",
    "    test_powers: Dict[str, int] = {}\n",
    "    all_ends = []\n",
    "\n",
    "    # Creates intervals.\n",
    "    for _, row in tests_data.iterrows():\n",
    "        name: str = row[\"Name\"]\n",
    "        operator: str = row[\"Operator\"]\n",
    "        test_time: int = row[\"TestTime\"]\n",
    "        average_power: int = row[\"AveragePower\"]\n",
    "        supply: str = operator_to_supply[operator]\n",
    "\n",
    "        start = model.new_int_var(0, horizon - test_time, f\"start_{name}\")\n",
    "        interval = model.new_fixed_size_interval_var(\n",
    "            start, test_time, f\"interval_{name}\"\n",
    "        )\n",
    "\n",
    "        # Bookkeeping.\n",
    "        test_starts[name] = start\n",
    "        test_durations[name] = test_time\n",
    "        test_powers[name] = average_power\n",
    "        test_supply[name] = supply\n",
    "        if supply not in tests_per_supply.keys():\n",
    "            tests_per_supply[supply] = ([], [])\n",
    "        tests_per_supply[supply][0].append(interval)\n",
    "        tests_per_supply[supply][1].append(average_power)\n",
    "        all_ends.append(start + test_time)\n",
    "\n",
    "    # Create supply cumulative constraints.\n",
    "    for supply, (intervals, demands) in tests_per_supply.items():\n",
    "        model.add_cumulative(intervals, demands, supply_to_max_power[supply])\n",
    "\n",
    "    # Objective.\n",
    "    makespan = model.new_int_var(0, horizon, \"makespan\")\n",
    "    for end in all_ends:\n",
    "        model.add(makespan >= end)\n",
    "    model.minimize(makespan)\n",
    "\n",
    "    # Solve model.\n",
    "    solver = cp_model.CpSolver()\n",
    "    if _PARAMS.value:\n",
    "        text_format.Parse(_PARAMS.value, solver.parameters)\n",
    "    status = solver.solve(model)\n",
    "\n",
    "    # Report solution.\n",
    "    if status == cp_model.OPTIMAL or status == cp_model.FEASIBLE:\n",
    "        print(f\"Makespan = {solver.value(makespan)}\")\n",
    "        for name, start in test_starts.items():\n",
    "            print(\n",
    "                f\"{name}: start:{solver.value(start)} duration:{test_durations[name]}\"\n",
    "                f\" power:{test_powers[name]} on supply {test_supply[name]}\"\n",
    "            )\n",
    "\n",
    "\n",
    "def main(argv: Sequence[str]) -> None:\n",
    "    \"\"\"Builds the data and solve the scheduling problem.\"\"\"\n",
    "    if len(argv) > 1:\n",
    "        raise app.UsageError(\"Too many command-line arguments.\")\n",
    "\n",
    "    tests_data, operators_data, supplies_data = build_data()\n",
    "    print(\"Tests data\")\n",
    "    print(tests_data)\n",
    "    print()\n",
    "    print(\"Operators data\")\n",
    "    print(operators_data)\n",
    "    print()\n",
    "    print(\"Supplies data\")\n",
    "    print(supplies_data)\n",
    "\n",
    "    solve(tests_data, operators_data, supplies_data)\n",
    "\n",
    "\n",
    "main()\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
